import torch
import os
from torch import nn
from PIL import Image
from torchvision import transforms
from models.encoders.psp import pSp
from models.stylegan2.model import Generator
from  models.encoders.psp_encoders import Encoder4Editing

# 全局变量用于存储初始化的模型
global_model = None

def process_image(image_path, latent_save_name):
    global global_model
    #获取图片名字
    base_path = "C:\\Users\\Dell\\PycharmProjects"

    # 使用 os.path.join 连接基础路径和相对路径
    image_path = os.path.join(base_path, image_path)
    print(image_path)
    image_name = os.path.basename(image_path)

    # Load the image
    image = Image.open(image_path).convert("RGB")
    #修改图片大小为1024*1024
    image = image.resize((1024, 1024), Image.BICUBIC)

    # Apply necessary transformations
    transform = transforms.Compose([
        transforms.Resize((512, 512)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])
    image = transform(image)
    image = image.unsqueeze(0)  # add batch dimension
    # Define the options for pSp
    device='cuda:1' if torch.cuda.is_available() else 'cpu'
    opts = {
        'encoder_type': 'Encoder4Editing',  # or 'Encoder4Editing', 'SingleStyleCodeEncoder'
        'stylegan_size': 1024,
        'return_latents': True,
        'stylegan_weights': 'stylegan2-ffhq-config-f.pt',
        'checkpoint_path': os.path.join(os.getcwd(), 'service/styleCLIP/e4e_ffhq_encode.pt'),
        'start_from_latent_avg': True,
        'device': device,
    }
    print(opts)
    # Check if the model has been initialized
    if global_model is None:
        # Instantiate pSp
        global_model = pSp(opts).to(device)
    # Convert the image to a latent code
    image = image.to(device)
    latents = global_model(image)
    # 计算新的形状
    new_shape = (latents.shape[0] * latents.shape[1], latents.shape[2])  # [4*18, 512]
    # 改变张量的形状
    latents2 = latents.contiguous().view(new_shape)
    # 计算第一个维度上的平均值，结果的形状为[1, 512]
    latents3 = latents2.mean(dim=0, keepdim=True)
    #保存变量
    latent_save_path = os.path.join('Parameter', latent_save_name)
    torch.save(latents3, latent_save_path)
    print(latents3,"保存完成,保存到到："+latent_save_path)
    return latent_save_path, image_name